# 本文件用于对电子云作图
# 由于是柱对称分子，仅画出轴向距离和径向距离两个参数的电子密度分布

from .molecules import *


def eval_gaussian(g: Gaussion_s, z: float, r: float) -> float:
    '''
    计算高斯函数
    参数：
        g: 1s 高斯轨道
        z: 轴向距离
        r: 径向距离
    '''
    return (2*g.alpha/np.pi)**(3/4) * np.exp(-g.alpha * ((z - g.center)**2 + r**2))


def eval_sto3g(s: STO3G, z: float, r: float) -> float:
    '''
    计算STO3G波函数
    参数：
        s: 1s STO3G轨道
        z: 轴向距离
        r: 径向距离
    '''
    total = 0.
    for i in range(3):
        total += s.d[i] * eval_gaussian(s.g[i], z, r)
    return total


def eval_cloud(basis: list, C: np.ndarray, z: float, r: float) -> float:
    '''
    计算电子云密度
    参数：
        basis: 基矢
        C:     波函数在基矢上的展开系数
        z:     轴向距离
        r:     径向距离
    '''
    assert len(basis) == len(C), "基矢维度和展开系数维度不一致"
    total = 0.
    N = len(basis)
    for i in range(N):
        for j in range(N):
            total += 2 * C[i] * C[j] * \
                eval_sto3g(basis[i], z, r) * eval_sto3g(basis[j], z, r)
    return total
